#include "slave.h"
#include "mpi_partitionings.h"
#include "task_list.h"
#include "defines.h"

#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <math.h>

void slaveMain( ConfigData *data ) {

  // Get information from the ConfigData structure passed in
  int rank            = data->mpi_rank;
  int num_procs       = data->mpi_procs;          // Used for number of strips/blocks
  PartType part_mode  = data->part_mode;
  
  int strip_height    = data->strip_height;       // Used for cyclical allocation.

  unsigned int height = data->height;    // How many pixels high to make the rendered image.
  unsigned int width  = data->width;     // How many pixels wide to make the rendered image.
  
  unsigned int part_size = ceil(sizeof(float) * width * height * 3);
  
  float *samples;       // Where we store the ray trace samples.
  float *send_arr;
  int* recv_arr;
  int* send_arr_int;
  int count, pixel;
  
  // Loop counter variables
  unsigned int x, y, i;
  int done = 0;
  
  int number_of_comps = 0;
  
  // Time variables
  double comp_start, comp_stop, comp_time;
  
  // Control variables
  unsigned int start_height, end_height;
  unsigned int start_width, end_width;
  
  int send_addr;
  int send_data_len;
  
  MPI_Status status;
  
  // Task list
  task_list* my_tasks;
  int is_first_task = 0;
  task_list* iter;

  // Give some information to the log
  printf ("(%d) Starting!\n", rank);  

  //////////////////////////////////////////////////////////////////////////////
  // Ray trace every point.
  //////////////////////////////////////////////////////////////////////////////
  
  switch (part_mode) {
  //////////////////////////////////////////////////////////////////////////////
    case PART_STATIC_STRIPS:
    
      samples = (float*) malloc( part_size );
      
      comp_start = MPI_Wtime();
      
      partition_static_strips_func (samples, data);
      
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;
            
      break;
      
  //////////////////////////////////////////////////////////////////////////////
    case PART_STATIC_BLOCKS:
      samples = (float*) malloc( part_size );

      comp_start = MPI_Wtime();

      partition_static_blocks_func (samples, data, &my_tasks);
      
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;
      
      break;

  //////////////////////////////////////////////////////////////////////////////
    case PART_STATIC_CYCLES:

      samples = (float*) malloc( part_size );
      comp_start = MPI_Wtime();
      
      partition_static_cycles_func (samples, data);
      
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;  

      break;
      
  //////////////////////////////////////////////////////////////////////////////
    case PART_DYNAMIC:
      samples = (float*) malloc( part_size );

      recv_arr = (int*) malloc (4 * sizeof(int));

      comp_start = MPI_Wtime();
      done = 0;
      while (!done) {

        // Send my rank to the master, requesting work
        MPI_Send (&rank, 1, MPI_INT, 0, TAG_DYN_REQ_WORK, MPI_COMM_WORLD);
      
        // Get an array with the pixels I need to take care of
        MPI_Recv (recv_arr, 4, MPI_INT, 0, TAG_DYN_GIVE_WORK, MPI_COMM_WORLD, &status);
        
        // Get info from send_arr
        start_height = recv_arr[0];
        start_width  = recv_arr[1];
        end_height   = recv_arr[2];
        end_width    = recv_arr[3];

        // Check for finish packet
        if ((int)start_height == -1 && (int)start_width == -1 && (int)end_height == -1 && (int)end_width == -1) {
          done = 1;
        } else {
        
          number_of_comps++;
        
          // Save task_list into local task_list
          if (is_first_task == 1) {
            is_first_task = 0;
            my_tasks = create_task_list(start_height, start_width, end_height, end_width);
          } else {
            my_tasks = add_task(my_tasks, start_height, start_width, end_height, end_width);
          }
          
          for (x = start_width; x <= end_width; x++) {
            for (y = start_height; y <= end_height; y++) {
              i = (y * width + x) * 3;
              shadePixel( samples + i, x, y, data );
            }
          }
          
        }
        
      }
      comp_stop = MPI_Wtime();
      comp_time = comp_stop - comp_start;  
      
      free (recv_arr);
      break;
    
  //////////////////////////////////////////////////////////////////////////////
  }

  printf ("(%d) Done with computations (%f seconds)\n", rank, comp_time);
  
  //////////////////////////////////////////////////////////////////////////////
  // Send data back to the master
  //////////////////////////////////////////////////////////////////////////////
   switch (part_mode) {
    case PART_STATIC_STRIPS:
      start_width = 0;
      end_width   = width - 1;
      
      start_height = rank * ceil(height/num_procs);
      
      // Take into account non-nicely-divisable height / num_proc
      end_height   = (rank == (num_procs-1)) ? height - 1 : (rank+1) * ceil(height/num_procs) - 1;
      
      send_addr     = ceil(start_height*3*width);
      send_data_len = 3*(end_height-start_height+1)*width;
      
      MPI_Send (&samples[send_addr], send_data_len, MPI_FLOAT, 0, 0, MPI_COMM_WORLD);
      
      break;
      
    case PART_STATIC_BLOCKS:

      // Malloc send array
      send_arr     = (float *) malloc (part_size);
      send_arr_int = (int *)   malloc (4 * sizeof(int));

      get_array(my_tasks, send_arr_int);
      start_height = send_arr_int[0];
      start_width  = send_arr_int[1];
      end_height   = send_arr_int[2];
      end_width    = send_arr_int[3];

      
      count = 0;
      for (x = start_width; x <= end_width; x++) {
        for (y = start_height; y <= end_height; y++) {
          // Check if y is outside range of height
          if (y > height) {
            break;
          }
          for (pixel = 0; pixel < 3; pixel++) {
            // Add to array
            send_arr[count++] = samples[(y * width + x) * 3 + pixel];
          }
        }
      }
      
      MPI_Send (send_arr, count, MPI_FLOAT, 0, 0, MPI_COMM_WORLD);
      free (send_arr);
      free (send_arr_int);
      
      break;

    case PART_STATIC_CYCLES:
    
      // Malloc send array
      send_arr = (float *) malloc (part_size);
      
      start_width = 0;
      end_width   = width - 1;

      // Need to put the samples into a single array and send over
      count = 0;
      for (x = start_width; x <= end_width; x++) {
        start_height = rank * strip_height;
        while (start_height < height) {
          for (y = start_height; y < start_height + strip_height; y++) {
            // Check if y is outside range of height
            if (y > height) {
              break;
            }
            for (pixel = 0; pixel < 3; pixel++) {
              // Add to array
              send_arr[count++] = samples[(y * width + x) * 3 + pixel];
            }
          }
          start_height = start_height + num_procs*strip_height;
        }
      }

      MPI_Send (send_arr, count, MPI_FLOAT, 0, 0, MPI_COMM_WORLD);
      
      free (send_arr);
      break;
      
    case PART_DYNAMIC:
        // Send the data back to the master
        send_arr_int = (int *)   malloc (4 * sizeof(int));
        send_arr     = (float *) malloc (part_size);
        int number_of_iters = 0;
        for (iter = my_tasks; iter != NULL; iter = iter->next) {
          if (number_of_iters >= number_of_comps ) {
            break;
          }
          
          get_array (iter, send_arr_int);
          start_height = send_arr_int[0];
          start_width  = send_arr_int[1];
          end_height   = send_arr_int[2];
          end_width    = send_arr_int[3];
          
          // Send array containing the dims of the work done
          MPI_Send(send_arr_int, 4, MPI_INT, 0, TAG_DYNAMIC_DATA_SEND1, MPI_COMM_WORLD);
          count = 0;
          for (x = start_width; x <= end_width; x++) {
            for (y = start_height; y <= end_height; y++) {
              for (pixel = 0; pixel < 3; pixel++) {
                // Add to array
                send_arr[count] = samples[(y * width + x) * 3 + pixel];
                count = count + 1;
              }
            }
          }
          MPI_Send(send_arr, count, MPI_FLOAT, 0, TAG_DYNAMIC_DATA_SEND2, MPI_COMM_WORLD);
          
          number_of_iters++;
        }
        free (send_arr);
        free (send_arr_int);
        printf("(%d) Data has been sent\n", rank);
      break;
  }
  
  // Send computation time to the master
  MPI_Send (&comp_time, 1, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
    
  free(samples);

  printf ("(%d) Exiting!\n", rank);  
}
